// Copyright 2024 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import {
  UseQueryOptions,
  UseQueryResult,
  useQueries,
} from '@tanstack/react-query';
import { isEqual } from 'lodash-es';
import { useCallback, useMemo, useRef, useState } from 'react';

import { Range, intersect, splitRange } from '@/generic_libs/tools/num_utils';

export const PENDING_RESULT: Omit<
  UseQueryResult<never, never>,
  'refetch' | 'remove' | 'promise'
> = Object.freeze({
  data: undefined,
  dataUpdatedAt: 0,
  error: null,
  errorUpdatedAt: 0,
  failureCount: 0,
  failureReason: null,
  errorUpdateCount: 0,
  isError: false,
  isFetched: false,
  isFetchedAfterMount: false,
  isFetching: true,
  isPending: true,
  isLoading: true,
  isLoadingError: false,
  isInitialLoading: true,
  isPaused: false,
  isPlaceholderData: false,
  isPreviousData: false,
  isRefetchError: false,
  isRefetching: false,
  isStale: false,
  isSuccess: false,
  status: 'pending',
  fetchStatus: 'idle',
});

export interface UseVirtualizedQueryOption<
  TQueryFnData,
  TDataItem,
  TError = unknown,
> {
  /**
   * The boundary of the range.  Only items with logical index within
   * `boundRange` will be queried. An `boundRange[1]` less than
   * `boundRange[0]` is treated the same as an `boundRange[1]` equals to
   * `boundRange[0]`.
   *
   * See `UseVirtualizedQueryResult.setRange` for details.
   */
  readonly rangeBoundary: Range;
  /**
   * The interval at which `[startBound, endBound)` is divided into chunks. When
   * the hook query items, it will only query one chunk in one query. Must be a
   * positive integer.
   *
   * For example, when `[startBound, endBound) = [103, 165)` and
   * `interval = 20`, the range will be divided in to the following chunks:
   * `[103, 120)`, `[120, 140)`, `[140, 160)`, `[160, 165)`.
   */
  readonly interval: number;
  /**
   * The initial range of items to be queried. See
   * `UseVirtualizedQueryResult.setRange` for details.
   */
  readonly initRange?: Range;
  /**
   * A function that generates a query that returns an array of items with
   * logical index `[start, end)`.
   *
   * `[start, end)` is guaranteed to be a single chunk obtained from dividing
   * `[startBound, endBound)` at multiples of `interval`. See `interval` and
   * `initDirection` for more details.
   */
  readonly genQuery: (
    start: number,
    end: number,
  ) => UseQueryOptions<TQueryFnData, TError, readonly TDataItem[]>;
  /**
   * The minimum duration a chunk needs to stay in the requested range before a
   * query is sent to query items in that chunk.
   *
   * The queries generated by the initial chunks requested by `initRange` are
   * always sent immediately.
   */
  readonly delayMs?: number;
}

export interface UseVirtualizedQueryResult<TDataItem, TError = unknown> {
  /**
   * Set the range to be queried. Only items with indices within the range will
   * be queried.
   *
   * Note that:
   *  * `[start, end)` is always expanded to `[alignedStart, alignedEnd)` where
   *    `[alignedStart, alignedEnd)` is the minimum superset of `[start, end)`
   *     and `alignedStart`/`alignedEnd` are both chunk boundaries (see
   *     `UseVirtualizedQueryOption.interval` for details).
   *  * only `setRange` calls that change `[alignedStart, alignedEnd)` will
   *    schedule an update to the queried chunks (and debounce the previously
   *    scheduled queries if `initDelayMs` is defined).
   *  * `[start, end)` should be a subset of `[startBound, endBound)`. If not,
   *    `[start, end)` is effectively shrunk to
   *    `[start, end) ∩ [startBound, endBound)`
   */
  setRange([start, end]: Range): void;
  /**
   * Get the query result for item at `index`. If index is not in
   * `[alignedStart, alignedEnd)`, the returned query result is always a
   * `PENDING_RESULT`.
   *
   * The range can be set using `setRange`.
   */
  get(
    index: number,
  ): Omit<UseQueryResult<TDataItem, TError>, 'refetch' | 'remove' | 'promise'>;
}

/**
 * A utility to divide a query into chunks that can be loaded/cached
 * individually.
 */
export function useVirtualizedQuery<TQueryFnData, TDataItem, TError = unknown>(
  opts: UseVirtualizedQueryOption<TQueryFnData, TDataItem, TError>,
): UseVirtualizedQueryResult<TDataItem, TError | Error> {
  const [startBound, endBound] = opts.rangeBoundary;
  const { ticks, base, getChunkIndexRange } = useMemo(() => {
    // Splits the boundary into multiple aligned chunks. So
    // 1. we can query one chunk at a time, and
    // 2. updating the chunk boundary will not cause all the query caches to be
    //    invalidated.
    const ticks = splitRange([startBound, endBound], opts.interval);

    // Pre-compute the base so the chunk index of an item index can be obtained
    // from `Math.floor((itemIndex - base) / opts.interval)`.
    const base = Math.floor(startBound / opts.interval) * opts.interval;

    /**
     * A utility function to help computing chunk index range of a given item
     * index range.
     */
    function getChunkIndexRange(itemIndexRange: Range): Range {
      const [start, end] = itemIndexRange;
      if (start >= end) {
        // An empty range.
        return [0, 0];
      }
      return [
        Math.max(Math.floor((start - base) / opts.interval), 0),
        Math.min(Math.ceil((end - base) / opts.interval), ticks.length - 1),
      ];
    }
    return { ticks, base, getChunkIndexRange };
  }, [startBound, endBound, opts.interval]);

  // The index range of items being requested currently.
  const rangeRef = useRef(opts.initRange || ([0, 0] as const));
  const chunkIndexRange = getChunkIndexRange(rangeRef.current);

  // The index range of items being requested `opts.delayMs` ago.
  // We only query items in `range ∩ settledRange`. This allows us to skip
  // loading items if they are scrolled away immediately.
  const settledRangeRef = useRef(rangeRef.current);
  const settledChunkIndexRange = getChunkIndexRange(settledRangeRef.current);

  // Used to trigger rerender.
  // We use this because updates to range and settled range do not necessarily
  // need to trigger refresh.
  const [_, setState] = useState({});

  const getChunkIndexRangeRef = useRef(getChunkIndexRange);
  getChunkIndexRangeRef.current = getChunkIndexRange;

  // Wrap the function in `useCallback` so ESLint can help us ensure that this
  // function has no dependency on local variables. Therefore it is safe to use
  // in a timeout.
  const scheduleQuery = useCallback((range: Range) => {
    const getChunkIndexRange = getChunkIndexRangeRef.current;
    const chunkIndexRange = getChunkIndexRange(rangeRef.current);

    // Compute the chunks that are being actively queried.
    const settledChunkIndexRange = getChunkIndexRange(settledRangeRef.current);
    const activeChunkIndexRange = intersect(
      settledChunkIndexRange,
      chunkIndexRange,
    );

    settledRangeRef.current = range;

    // Compute the chunks that will be actively queried.
    const nextSettledChunkIndexRange = getChunkIndexRange(
      settledRangeRef.current,
    );
    const nextActiveChunkIndexRange = intersect(
      nextSettledChunkIndexRange,
      chunkIndexRange,
    );

    // We only need to trigger re-render to update which query is enabled if
    // there's an update to active chunks.
    if (isEqual(activeChunkIndexRange, nextActiveChunkIndexRange)) {
      return;
    }
    setState({});
  }, []);

  const queries = Array(chunkIndexRange[1] - chunkIndexRange[0])
    .fill(undefined)
    .map((_, i) => {
      const chunkIndex = chunkIndexRange[0] + i;
      const queryOpts = opts.genQuery(ticks[chunkIndex], ticks[chunkIndex + 1]);
      return {
        ...queryOpts,
        enabled:
          (queryOpts.enabled ?? true) &&
          // Do not send query for chunks that are newly requested. This allows
          // us to skip loading them that are scrolled away immediately.
          chunkIndex >= settledChunkIndexRange[0] &&
          chunkIndex < settledChunkIndexRange[1],
      };
    });
  const queryResults = useQueries({ queries });

  return {
    setRange(range) {
      const prevChunkIndexRange = getChunkIndexRange(rangeRef.current);
      rangeRef.current = range;
      const newChunkIndexRange = getChunkIndexRange(rangeRef.current);

      // We only need to trigger re-render to update the queries if there's an
      // update to the queried chunks.
      if (!isEqual(newChunkIndexRange, prevChunkIndexRange)) {
        setState({});
      }

      if (opts.delayMs) {
        setTimeout(() => scheduleQuery(range), opts.delayMs);
      } else {
        scheduleQuery(range);
      }
    },
    get(index) {
      if (
        index < ticks[chunkIndexRange[0]] ||
        index >= ticks[chunkIndexRange[1]]
      ) {
        return PENDING_RESULT;
      }
      const chunkIndex = Math.floor((index - base) / opts.interval);
      const pageIndex = chunkIndex - chunkIndexRange[0];
      const itemIndex = index - ticks[chunkIndex];
      const page = queryResults[pageIndex];
      return {
        ...page,
        data: page.data?.[itemIndex],
      };
    },
  };
}
